{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "551af157-d808-4aba-be70-2cad33af8b10",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2024-05-20T08:19:40.656619Z",
     "iopub.status.busy": "2024-05-20T08:19:40.656278Z",
     "iopub.status.idle": "2024-05-20T08:19:43.651564Z",
     "shell.execute_reply": "2024-05-20T08:19:43.651057Z",
     "shell.execute_reply.started": "2024-05-20T08:19:40.656598Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from lorahub.algorithm import *\n",
    "import json\n",
    "\n",
    "import sys, os \n",
    "def get_ds_list(root):\n",
    "    task = [\"ner\", \"re\", \"ee\"]\n",
    "\n",
    "    ner = []\n",
    "    re = []\n",
    "    ee = []\n",
    "\n",
    "    for t in task:\n",
    "        path = root + t + \"/\"\n",
    "        for ds in os.listdir(path):\n",
    "            if t == \"ner\":\n",
    "                ner.append(path + ds)\n",
    "            elif t == \"re\":\n",
    "                re.append(path + ds)\n",
    "            elif t == \"ee\":\n",
    "                ee.append(path + ds)\n",
    "    return ner, re, ee\n",
    "\n",
    "def get_lora_list(root, task=\"all\"):\n",
    "    ner, re, ee = get_ds_list(root)\n",
    "\n",
    "    if task == \"ner\":\n",
    "        return ner\n",
    "    elif task == \"re\":\n",
    "        return re\n",
    "    elif task == \"ee\":\n",
    "        return ee\n",
    "    else:\n",
    "        return ner + re + ee\n",
    "\n",
    "# print(get_lora_list(\"/mnt/workspace/save/t5_xl/\"))\n",
    "\n",
    "\n",
    "\n",
    "def get_examples_for_learning(p):\n",
    "    res = []\n",
    "    for path in p:\n",
    "        with open(path, 'r') as f:\n",
    "              js = json.load(f)\n",
    "        res.extend(js)\n",
    "    return res\n",
    "\n",
    "\n",
    "def get_examples_for_inference(p):\n",
    "    res = []\n",
    "    for path in p:\n",
    "        with open(path, 'r') as f:\n",
    "              js = json.load(f)\n",
    "        res.extend(js)\n",
    "    return res\n",
    "\n",
    "\n",
    "def get_lora_module_list(task):\n",
    "    pefix = \"/mnt/save/t_xl/\"\n",
    "\n",
    "    all_loras = ['re_task/duie_re', 'ee_task/ace05_tuple_ee', 'ee_task/casie_tuple_ee', 'ee_task/duee_tuple_ee', 'ee_task/genia_tuple_ee',\n",
    "                 'ee_task/phee_tuple_ee', 'ner_task/ace_ner', 'ner_task/cnerta_ner', 'ner_task/conll_ner',\n",
    "                 'ner_task/multinerd_ner', 're_task/conll04_re', 're_task/gids_re',\n",
    "                 're_task/nyt11_re']\n",
    "\n",
    "    ner_loras = [\"ace05\",  \"conll03\",  \"mit-movie\",  \"mit-restaurant\",  \"ontonotes\"  \"wnut17\"]\n",
    "\n",
    "    re_loras = ['re_task/duie_re', 're_task/conll04_re',  're_task/gids_re', 're_task/nyt11_re']\n",
    "\n",
    "    ee_loras = ['ee_task/duee_tuple_ee', 'ee_task/ace05_tuple_ee', 'ee_task/casie_tuple_ee', 'ee_task/genia_tuple_ee',\n",
    "                'ee_task/phee_tuple_ee']\n",
    "\n",
    "    if task == 'ner':\n",
    "        p = ner_loras\n",
    "    elif task == 're':\n",
    "        p = re_loras\n",
    "    elif task == 'ee':\n",
    "        p = ee_loras\n",
    "    elif task == 'all':\n",
    "        p = all_loras\n",
    "    else:\n",
    "        raise ValueError(\"task must be one of 'ner', 're', 'ee', 'all'\")\n",
    "\n",
    "    return [pefix + x for x in p]\n",
    "\n",
    "def perform_few_shot_learning(target_ds,  lora_list, example_inputs, examples_outputs):\n",
    "    models = []\n",
    "    for l in lora_list:\n",
    "        if l.endswith(\".ipynb_checkpoints\"): continue\n",
    "        if l.endswith(\"_f\"):\n",
    "            if target_ds in l :\n",
    "                models.append(l)\n",
    "        else:\n",
    "            if target_ds in l :\n",
    "                continue\n",
    "            else :\n",
    "                models.append(l)\n",
    "\n",
    "    print(models)\n",
    "    # return \n",
    "\n",
    "    module_weights, model, tokenizer = lorahub_learning(lora_module_list=models,\n",
    "                                                        example_inputs=example_inputs,\n",
    "                                                        example_outputs=examples_outputs,\n",
    "                                                        max_inference_step=40,\n",
    "                                                        batch_size=32,\n",
    "                                                        )\n",
    "\n",
    "    for m, w in zip(models, module_weights):\n",
    "        print(f\"{m} : \\t\\t\\t{w}\")\n",
    "\n",
    "    return model, tokenizer, module_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed9744a8-d499-4921-9a8d-f1d3fc5af0af",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-05-20T08:20:12.116431Z",
     "iopub.status.busy": "2024-05-20T08:20:12.115949Z",
     "iopub.status.idle": "2024-05-20T08:20:44.658750Z",
     "shell.execute_reply": "2024-05-20T08:20:44.658274Z",
     "shell.execute_reply.started": "2024-05-20T08:20:12.116402Z"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/mnt/workspace/save/t5_xl/ner/wnut17', '/mnt/workspace/save/t5_xl/ner/mit-restaurant', '/mnt/workspace/save/t5_xl/ner/ace05', '/mnt/workspace/save/t5_xl/ner/ontonotes', '/mnt/workspace/save/t5_xl/ner/conll03', '/mnt/workspace/save/t5_xl/ner/mit-movie', '/mnt/workspace/save/t5_xl/re/scierc', '/mnt/workspace/save/t5_xl/re/nyt', '/mnt/workspace/save/t5_xl/re/nyt11', '/mnt/workspace/save/t5_xl/re/gids', '/mnt/workspace/save/t5_xl/ee/phee', '/mnt/workspace/save/t5_xl/ee/casie', '/mnt/workspace/save/t5_xl/ee/ace05', '/mnt/workspace/save/t5_xl/ee/genia']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00,  1.55it/s]\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'T5Tokenizer'. \n",
      "The class this function is called from is 'PreTrainedTokenizerFast'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Begin to load lora modules\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]/usr/local/lib/python3.10/dist-packages/peft/utils/save_and_load.py:195: UserWarning: Could not find a config file in /mnt/workspace/model_save/t5_xl - will assume that the vocabulary was not modified.\n",
      "  warnings.warn(\n",
      "1it [00:00,  5.22it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ner/wnut17 ...\n",
      "> Loading /mnt/workspace/save/t5_xl/ner/mit-restaurant ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [00:00,  4.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ner/ace05 ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "3it [00:00,  4.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ner/ontonotes ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4it [00:00,  4.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ner/conll03 ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [00:01,  4.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ner/mit-movie ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "6it [00:01,  4.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/re/scierc ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "7it [00:01,  4.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/re/nyt ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "9it [00:01,  4.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/re/nyt11 ...\n",
      "> Loading /mnt/workspace/save/t5_xl/re/gids ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "10it [00:02,  4.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ee/phee ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "11it [00:02,  4.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ee/casie ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "12it [00:02,  4.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ee/ace05 ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "13it [00:02,  3.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading /mnt/workspace/save/t5_xl/ee/genia ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "14it [00:03,  4.35it/s]\n",
      "Running tokenizer on dataset: 100%|██████████| 25/25 [00:00<00:00, 2441.56 examples/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['input_ids', 'attention_mask', 'labels'],\n",
      "    num_rows: 25\n",
      "})\n",
      "> Begin to perform gradient-free optimization ...\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:01,  1.07s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.08125\n",
      "39 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.06397321428571429\n",
      "38 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.04669642857142857\n",
      "37 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03973214285714286\n",
      "36 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.04776785714285715\n",
      "35 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.05151785714285714\n",
      "34 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.046830357142857146\n",
      "33 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03620535714285714\n",
      "32 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03174107142857143\n",
      "31 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.037589285714285714\n",
      "30 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03258928571428571\n",
      "29 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.034308035714285715\n",
      "28 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03837053571428571\n",
      "27 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03446428571428571\n",
      "26 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.05977678571428571\n",
      "25 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.04612507505436695\n",
      "24 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03040990422670417\n",
      "23 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03057815557070047\n",
      "22 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.031172577663970774\n",
      "21 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.03169877668755576\n",
      "20 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.034594628452776235\n",
      "19 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.027692227825109175\n",
      "18 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.027617510034450413\n",
      "17 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025084482238313888\n",
      "16 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.026292983058101814\n",
      "15 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.02504580965498909\n",
      "14 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025331870206794248\n",
      "13 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.02550694555135081\n",
      "12 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.029004716879435982\n",
      "11 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025342179746454217\n",
      "10 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.02553516434576858\n",
      "9 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025235677477432435\n",
      "8 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.026584508781339444\n",
      "7 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.02543237389643941\n",
      "6 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.02534751290890694\n",
      "5 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.026266322580420817\n",
      "4 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.02507751473339414\n",
      "3 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025231136391890834\n",
      "2 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025213258591817535\n",
      "1 remaining budget and 0 running jobs\n",
      "Launching 1 jobs with new suggestions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:00,  2.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating fitness with value 0.025848896788699085\n",
      "0 remaining budget and 0 running jobs\n",
      "/mnt/workspace/save/t5_xl/ner/wnut17 : \t\t\t-0.03661731514829045\n",
      "/mnt/workspace/save/t5_xl/ner/mit-restaurant : \t\t\t0.5260319279442587\n",
      "/mnt/workspace/save/t5_xl/ner/ace05 : \t\t\t0.19257240069681758\n",
      "/mnt/workspace/save/t5_xl/ner/ontonotes : \t\t\t0.0019348365996101767\n",
      "/mnt/workspace/save/t5_xl/ner/conll03 : \t\t\t-0.055195335022532516\n",
      "/mnt/workspace/save/t5_xl/ner/mit-movie : \t\t\t0.05720831059209082\n",
      "/mnt/workspace/save/t5_xl/re/scierc : \t\t\t0.3740141117433971\n",
      "/mnt/workspace/save/t5_xl/re/nyt : \t\t\t0.5651506992792246\n",
      "/mnt/workspace/save/t5_xl/re/nyt11 : \t\t\t-0.17908635807606854\n",
      "/mnt/workspace/save/t5_xl/re/gids : \t\t\t-0.012378632863052576\n",
      "/mnt/workspace/save/t5_xl/ee/phee : \t\t\t-0.03746165208555413\n",
      "/mnt/workspace/save/t5_xl/ee/casie : \t\t\t-0.014056523078421285\n",
      "/mnt/workspace/save/t5_xl/ee/ace05 : \t\t\t-0.03974192656032694\n",
      "/mnt/workspace/save/t5_xl/ee/genia : \t\t\t-0.02137667370729959\n"
     ]
    }
   ],
   "source": [
    "modules = get_lora_list(\"/mnt/workspace/save/t5_xl/\",\"all\")\n",
    "\n",
    "modules = [m for m in modules if not m.endswith(\".ipynb_checkpoints\") and not m.endswith(\"_f\")]\n",
    "\n",
    "example_inputs, examples_outputs = [], []\n",
    "with open(\"/mnt/workspace/data/RE/conll04_RE/5_shot_train.json\", \"r\", encoding=\"utf-8\") as f:\n",
    "        js = json.load(f)\n",
    "for example in js:\n",
    "    example_inputs.append(example[\"prompt\"])\n",
    "    examples_outputs.append(example[\"response\"])\n",
    "\n",
    "model, tokenizer, module_weights = perform_few_shot_learning(\"re/conll04\",modules, example_inputs, examples_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d4e756ed-716a-4af4-8735-b0dc73ad3c2a",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-05-20T08:20:50.593197Z",
     "iopub.status.busy": "2024-05-20T08:20:50.592843Z",
     "iopub.status.idle": "2024-05-20T08:21:04.757964Z",
     "shell.execute_reply": "2024-05-20T08:21:04.757307Z",
     "shell.execute_reply.started": "2024-05-20T08:20:50.593177Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running tokenizer on dataset: 100%|██████████| 288/288 [00:00<00:00, 3412.45 examples/s]\n",
      "  0%|          | 0/9 [00:00<?, ?it/s]/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py:2707: UserWarning: `max_length` is ignored when `padding`=`True` and there is no truncation strategy. To pad to max length, use `padding='max_length'`.\n",
      "  warnings.warn(\n",
      "100%|██████████| 9/9 [00:14<00:00,  1.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "task accuracy: {'precision': 0.24293785310734464, 'recall': 0.22395833333333334, 'f1': 0.23306233062330622}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "example_inputs, examples_outputs = [], []\n",
    "with open(\"/mnt/workspace/data/RE/conll04_RE/test.json\", \"r\", encoding=\"utf-8\") as f:\n",
    "        js = json.load(f)\n",
    "for example in js:\n",
    "    example_inputs.append(example[\"prompt\"])\n",
    "    examples_outputs.append(example[\"response\"])\n",
    "\n",
    "example_predictions1, perf = assess_task_performance(example_inputs=example_inputs,\n",
    "                                                    model_or_name_path=model,\n",
    "                                                    tokenizer_or_tokenizer_path=tokenizer,\n",
    "                                                    batch_size=32,\n",
    "                                                    # can set as None if you do not have the ground truth\n",
    "                                                    example_outputs=examples_outputs,\n",
    "                                                    task = \"re\"\n",
    "                                                    )\n",
    "# print(\"example_predictions:\", example_predictions)\n",
    "print(\"task accuracy:\", perf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "717d27f7-f293-4eb9-8904-7ab2dee5f518",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-05-20T08:21:13.691607Z",
     "iopub.status.busy": "2024-05-20T08:21:13.691261Z",
     "iopub.status.idle": "2024-05-20T08:21:13.699793Z",
     "shell.execute_reply": "2024-05-20T08:21:13.699132Z",
     "shell.execute_reply.started": "2024-05-20T08:21:13.691588Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'precision': 0.32526881720430106, 'recall': 0.2972972972972973, 'f1': 0.31065468549422337}\n"
     ]
    }
   ],
   "source": [
    "def ee_calculate_metrics(predicted: List[str], actual: List[str]):\n",
    "    def ee_parse_tuple_string(tuple_string: str) -> List[Tuple[str, str, str]]:\n",
    "        tuple_string = tuple_string.replace('，',',').replace('（','(').replace('）', ')').replace(\"：\", \":\").replace(\" \", \"\").replace(\"the\", \"\")\n",
    "        pattern = re.compile(r'\\(([^,]*),? ?([^,]*)?,? ?([^)]*)?\\)')\n",
    "        tuples = pattern.findall(tuple_string)\n",
    "        tuples = set(('EMPTY', 'EMPTY', 'EMPTY') if t == ('', '', '') else t for t in tuples)\n",
    "        return tuples\n",
    "\n",
    "    tp = 0\n",
    "    fp = 0\n",
    "    fn = 0\n",
    "\n",
    "    for p, a in zip(predicted, actual):\n",
    "        predicted_tuples = ee_parse_tuple_string(p)\n",
    "        actual_tuples = ee_parse_tuple_string(a)\n",
    "\n",
    "        tp_temp = 0\n",
    "        for pt in predicted_tuples:\n",
    "            for at in actual_tuples:\n",
    "                if (pt[0] in at[0] and pt[1] in at[1] and pt[2] in at[2]) or (at[0] in pt[0] and at[1] in pt[1] and at[2] in pt[2]):\n",
    "                    tp_temp += 1\n",
    "                    break\n",
    "\n",
    "        tp += tp_temp\n",
    "        fp += len(predicted_tuples) - tp_temp\n",
    "        fn += len(actual_tuples) - tp_temp\n",
    "\n",
    "    precision = tp / (tp + fp) if (tp + fp) > 0 else 1\n",
    "    recall = tp / (tp + fn) if (tp + fn) > 0 else 1\n",
    "    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 1\n",
    "    \n",
    "    # print(tp, fp, fn)\n",
    "\n",
    "    return {\"precision\": precision, \"recall\": recall, \"f1\": f1}\n",
    "print(ee_calculate_metrics(example_predictions1, examples_outputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f348a6a-08fe-493d-b216-bd0fce37ca01",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
